package nmap

import (
	"bufio"
	"bytes"
	"encoding/hex"
	"fmt"
	logx "gitee.com/wudicaidou/menciis-logx"
	"github.com/dlclark/regexp2"
	"github.com/pkg/errors"
	"net/url"
	"strconv"
	"strings"
	"time"
	"unicode"
)

const (
	DirectiveExclude    = "exclude"
	DirectiveProbe      = "probe"
	DirectiveMatch      = "match"
	DirectiveSoftMatch  = "softmatch"
	DirectivePorts      = "ports"
	DirectiveSSLPorts   = "sslports"
	DirectiveTotalWait  = "totalwaitms"
	DirectiveTcpWrapped = "tcpwrappedms"
	DirectiveRarity     = "rarity"
	DirectiveFallback   = "fallback"
)

const DefaultMatchTimeout = time.Second * 1

const ProbeNULL = "NULL"

type ProbesParser struct {
	index  int
	result *ProbeParseResult
}

func NewProbesParser() *ProbesParser {
	return &ProbesParser{result: &ProbeParseResult{}}
}

const tplErrProbeNil = "probe can't be nil, %v"

func (p *ProbesParser) Parse(raw []byte) error {
	scanner := bufio.NewScanner(bytes.NewReader(raw))
	scanner.Split(bufio.ScanLines)
	nmapProbes := p.result

	var currentProbe *Probe
	for {
		directive, rest, err := p.next(scanner)
		if err != nil {
			return err
		}
		if directive == "" {
			//log.Debug("parse done")
			break
		}
		directive = strings.ToLower(directive)
		switch directive {
		case DirectiveExclude:
			ports, err := ParseRangePort(string(rest), TCP)
			if err != nil {
				logx.Errorf("parse exclude directive error, %s", err)
				continue
			}
			if nmapProbes.ExcludePorts != nil {
				return fmt.Errorf("file format error, you should upgrade to []ExcludePorts")
			}
			nmapProbes.ExcludePorts = ports
		case DirectiveFallback:
			if currentProbe == nil {
				return fmt.Errorf(tplErrProbeNil, directive+string(rest))
			}
			currentProbe.Fallback = p.parseFallback(rest)
		case DirectiveMatch, DirectiveSoftMatch:
			match, err := p.parseMatch(rest, directive == DirectiveSoftMatch)
			if err != nil {
				logx.Errorf("parse Match directive error, %s", err)
				continue
			}
			if currentProbe == nil {
				return fmt.Errorf(tplErrProbeNil, directive+string(rest))
			}
			currentProbe.Matches = append(currentProbe.Matches, match)
		case DirectivePorts, DirectiveSSLPorts:
			if currentProbe == nil {
				return fmt.Errorf(tplErrProbeNil, directive+string(rest))
			}
			ports, err := ParseRangePort(string(rest), currentProbe.Proto)
			if err != nil {
				logx.Errorf("parse ports/sslports error, %s", err)
				continue
			}
			if directive == DirectivePorts {
				currentProbe.Ports = ports
			} else {
				currentProbe.SSLPorts = ports
			}
		case DirectiveProbe:
			if currentProbe != nil {
				nmapProbes.Probes = append(nmapProbes.Probes, currentProbe)
			}
			probe, err := p.parseProbe(rest)
			if err != nil {
				// Probe fail 应该直接返回了，不然后面没法解析了
				return err
			}
			currentProbe = probe
		case DirectiveRarity:
			rarity, err := p.parseInt(rest)
			if err != nil {
				return err
			}
			if currentProbe == nil {
				return fmt.Errorf(tplErrProbeNil, directive+string(rest))
			}
			currentProbe.Rarity = rarity
		case DirectiveTcpWrapped:
			_, err := p.parseInt(rest)
			if err != nil {
				return err
			}
			if currentProbe == nil {
				return fmt.Errorf(tplErrProbeNil, directive+string(rest))
			}
			//log.Debug("got wrapped time ", wrapped)
		case DirectiveTotalWait:
			wait, err := p.parseInt(rest)
			if err != nil {
				return err
			}
			if currentProbe == nil {
				return fmt.Errorf(tplErrProbeNil, directive+string(rest))
			}
			currentProbe.WaitTimeout = time.Duration(wait) * time.Millisecond
		default:
			logx.Println("unsupported directive ", directive)
		}
	}

	// 最后一个
	if currentProbe != nil {
		nmapProbes.Probes = append(nmapProbes.Probes, currentProbe)
	}
	if nmapProbes.ExcludePorts == nil {
		nmapProbes.ExcludePorts = &RangePort{}
	}
	return nil
}

func (p *ProbesParser) Result() *ProbeParseResult {
	return p.result
}

func (p *ProbesParser) next(scanner *bufio.Scanner) (directive string, rest []byte, err error) {
	for {
		ok := scanner.Scan()
		if !ok {
			return "", nil, scanner.Err()
		}
		line := scanner.Bytes()
		line = bytes.TrimSpace(line)
		if len(line) == 0 || line[0] == '#' {
			continue
		}

		parts := bytes.SplitN(line, []byte{' '}, 2)
		if len(parts) != 2 {
			return "", nil, fmt.Errorf("can't split into two parts, %s", string(line))
		}
		directive, rest = string(parts[0]), parts[1]
		rest = bytes.TrimLeftFunc(rest, unicode.IsSpace)
		return directive, rest, nil
	}
}

func (p *ProbesParser) parseProbe(rest []byte) (*Probe, error) {
	ret := &Probe{
		Index: p.index,
	}
	p.index++
	parts := bytes.SplitN(rest, []byte{' '}, 3)
	if len(parts) != 3 {
		return nil, fmt.Errorf("length of blocks is not 4, %s", rest)
	}

	protoRaw, name, payload := string(parts[0]), string(parts[1]), parts[2]
	ret.Name = name
	switch TransportProto(protoRaw) {
	case TCP:
		ret.Proto = TCP
	case UDP:
		ret.Proto = UDP
	default:
		return nil, fmt.Errorf("proto error, (only TCP/UDP) support %s", protoRaw)
	}

	if len(payload) < 3 {
		return nil, fmt.Errorf("invalid payload length")
	}

	sepStart, sepEnd := payload[1], payload[len(payload)-1]

	if payload[0] != 'q' || sepStart != sepEnd {
		fmt.Println(string(payload), payload[0], payload[1], len(payload))
		return nil, errors.Errorf("invalid payload format: %s", rest)
	}

	realPayload := payload[2 : len(payload)-1]
	// NULL Probe
	if len(realPayload) == 0 {
		return ret, nil
	}

	var err error
	realPayload, err = UnquoteCStyleString(realPayload)
	if err != nil {
		return nil, fmt.Errorf("unquote payload[%v] failed: %s", payload, err)
	}
	ret.Payload = realPayload
	return ret, nil
}

const (
	matchStateNameStart = iota
	matchStateFieldStart
	matchStateInFieldName
	matchStateFieldDelimiterStart
	matchStateInFieldData
	matchStateFieldOption

	delimiter = ' '
)

func (p *ProbesParser) parseMatch(rest []byte, soft bool) (*Match, error) {
	m := &Match{
		Soft: soft,
	}
	state := matchStateNameStart
	var currentField string
	var currentValue string
	var currentDelimiter byte
	var buf bytes.Buffer

	for _, c := range rest {
		switch state {
		case matchStateNameStart:
			if c == delimiter {
				state = matchStateFieldStart
				m.ServiceName = buf.String()
				buf.Reset()
				continue
			}
			buf.WriteByte(c)
		case matchStateFieldStart:
			if c == delimiter {
				continue
			}
			switch c {
			case 'm', 'p', 'v', 'i', 'h', 'o', 'd':
				currentField = string(c)
				state = matchStateFieldDelimiterStart
			case 'c':
				// cpe
				buf.WriteByte(c)
				state = matchStateInFieldName
			default:
				return nil, fmt.Errorf("unexpected field start char, %c", c)
			}

		case matchStateInFieldName:
			switch c {
			case 'p', 'e':
				buf.WriteByte(c)
			case ':':
				currentField = buf.String()
				buf.Reset()
				state = matchStateFieldDelimiterStart
			default:
				return nil, fmt.Errorf("expected in field char, %c", c)
			}
		case matchStateFieldDelimiterStart:
			currentDelimiter = c
			state = matchStateInFieldData
		case matchStateInFieldData:
			if c == currentDelimiter {
				currentValue = buf.String()
				buf.Reset()
				state = matchStateFieldOption
				continue
			}
			buf.WriteByte(c)
		case matchStateFieldOption:
			if c == delimiter {
				if err := p.setFieldData(currentField, currentValue, buf.String(), m); err != nil {
					return nil, err
				}
				buf.Reset()
				state = matchStateFieldStart
				continue
			}
			buf.WriteByte(c)
		default:
			return nil, fmt.Errorf("unknow state, %d", state)
		}
	}
	// 写入最后一个 field, 不方便在状态机内直接实现
	if err := p.setFieldData(currentField, currentValue, buf.String(), m); err != nil {
		return nil, err
	}
	m.Rule.MatchTimeout = DefaultMatchTimeout
	return m, nil
}

func (p *ProbesParser) setFieldData(field, value, option string, m *Match) error {
	var err error
	if field != "m" {
		value, err = url.QueryUnescape(value)
		if err != nil {
			logx.Error(err)
			return err
		}
	}
	switch field {
	case "m":
		var opts regexp2.RegexOptions
		for _, opt := range option {
			if opt == 'i' {
				opts |= regexp2.IgnoreCase
			} else if opt == 's' {
				opts |= regexp2.Singleline
			} else {
				return fmt.Errorf("new regexp option found, please support it, option: %s, value: %s", option, value)
			}
		}
		value = strings.ReplaceAll(value, "\\0", "\x00")
		m.Rule = regexp2.MustCompile(value, opts)
	case "p":
		m.Product = value
	case "v":
		m.Version = value
	case "i":
		m.Info = value
	case "h":
		m.Hostname = value
	case "o":
		m.OS = value
	case "d":
		m.DeviceType = value
	case "cpe":
		m.CPE = append(m.CPE, fmt.Sprintf("cpe:/%s", value))
	default:
		return fmt.Errorf("unknow field type, %s", field)
	}
	return nil
}

func (p *ProbesParser) parseFallback(rest []byte) []string {
	parts := bytes.Split(rest, []byte{','})
	ret := make([]string, len(parts))
	for i, p := range parts {
		ret[i] = string(p)
	}
	return ret
}

func (p *ProbesParser) parseInt(rest []byte) (int, error) {
	return strconv.Atoi(string(rest))
}

func UnquoteCStyleString(raw []byte) ([]byte, error) {
	const (
		stateDefault = iota
		stateEscape
		stateHex
	)

	var state = stateDefault
	var buf bytes.Buffer
	hexBuffer := make([]byte, 2)
	var i int

	for _, c := range raw {
		if c == '\\' && state != stateEscape {
			state = stateEscape
			continue
		}

		switch state {
		case stateHex:
			if i == 0 {
				hexBuffer[i] = c
				i++
				continue
			}
			hexBuffer[1] = c
			// it's safe to use self as dst and src at the same time
			n, err := hex.Decode(hexBuffer, hexBuffer)
			if n != 1 || err != nil {
				return nil, fmt.Errorf("parse hex buffer[\\x%s] failed: %s", hexBuffer, err)
			}
			buf.WriteByte(hexBuffer[0])
			state = stateDefault
			i = 0

		case stateEscape:
			switch c {
			case '0':
				buf.WriteByte(0x00)
			case '\\':
				buf.WriteByte('\\')
			case 'a':
				buf.WriteByte('\a')
			case 'b':
				buf.WriteByte('\b')
			case 'f':
				buf.WriteByte('\f')
			case 'n':
				buf.WriteByte('\n')
			case 'r':
				buf.WriteByte('\r')
			case 't':
				buf.WriteByte('\t')
			case 'v':
				buf.WriteByte('\v')
			case 'x':
				state = stateHex
				continue
			}
			state = stateDefault
		default:
			buf.WriteByte(c)
		}

	}
	return buf.Bytes(), nil
}

func ParserMatchLine(raw string) (*Match, error) {
	parser := &ProbesParser{}
	if strings.HasPrefix(raw, "match ") {
		return parser.parseMatch([]byte(strings.TrimPrefix(raw, "match ")), false)
	} else if strings.HasPrefix(raw, "softmatch ") {
		return parser.parseMatch([]byte(strings.TrimPrefix(raw, "softmatch ")), true)
	} else {
		return nil, fmt.Errorf("unknow directive when parsing %s", raw)
	}
}
